from penalty_function import makeExactOne, gridSweep
from collections import defaultdict
from config import *

# 假设没有时间切面
# 总的时间就是
# queue_size 
all_nodes = list(range(node_number))

def NodeQueueSize(node, place):  # node的queue size的第place位
    return ('queue_size', node, place)

def NodeChosed(node, target_node):   # node的选择了target node执行文件
    return ('chose_node', node, target_node)

def NodeDataSize(node, place):  # node的要执行的函数的文件大小的第place位
    return ('data_size', node, place)

Q = defaultdict(float)

for node in all_nodes:
    data_size_qubits = [NodeDataSize(node, i) for i in range(max_data_qubit_num)]
    
    target_node_qubits = [NodeChosed(node, target_node) for target_node in all_nodes]
    empty_target_node_qubit = NodeChosed(node, empty_node)
    

    queue_size_qubits = [NodeQueueSize(node, i) for i in range(max_queue_qubit_num)]

    for target_node in all_nodes:
        target_node_qubit = NodeChosed(node, target_node)
        target_queue_size_qubits = [NodeQueueSize(target_node, i) for i in range(max_queue_qubit_num)]
        for place, target_queue_size_qubit in enumerate(target_queue_size_qubits):
            Q[(target_node_qubit, target_queue_size_qubit)] += avg_timecost * (2 ** place)  #如果发给target_node需要等待前面queue完成的时间
        
        if target_node != node:
            for place, data_size_qubit in enumerate(data_size_qubits):
                Q[(target_node_qubit, data_size_qubit)] += (2 ** place) / network_speed  #如果发给target_node网络传输需要花费的时间


    # 冲突导致的总时长增加
    call_node_qubits = [NodeChosed(call_node, target_node)  for call_node in all_nodes]
    # call_number * (call_number - 1) / 2 * avg_cost
    for call_node_qubit1, call_node_qubit2 in gridSweep(call_node_qubits):
        Q[(call_node_qubit1, call_node_qubit2)] += avg_timecost

    # 加下来是损失函数
    makeExactOne(Q, target_node_qubits + [empty_target_node_qubit],  100)   # 一次只能选一个bide执行
    # 如果选择的是empty node，data size应该为0
    for data_size_qubit in data_size_qubits:
        Q[(empty_target_node_qubit, data_size_qubit)] += 100

# print(Q)
from neal import SimulatedAnnealingSampler # cpu
import dimod


class Case():
    def __init__(self):
        self.chosed_nodes = None
        self.data_sizes = None
        self.queue_sizes = None 
        self.count = 0
        return

def filterCases(func, cases):
    return [case for case in cases if func(case)]

def analyzeCase(response):

        # def __hash__(self):
        #     return hash()
    all_cases = []

    for sample, energy, num_occurrences in response.data(['sample', 'energy', 'num_occurrences'])[:10]:
        data_sizes = []
        queue_sizes = []
        chosed_nodes = []
        for node in all_nodes:
            data_size_qubits = [NodeDataSize(node, i) for i in range(max_data_qubit_num)]
            target_node_qubits = [NodeChosed(node, target_node) for target_node in all_nodes]
            empty_target_node_qubit = NodeChosed(node, empty_node)
            queue_size_qubits = [NodeQueueSize(node, i) for i in range(max_queue_qubit_num)]

            data_size = 0
            for place, data_size_qubit in enumerate(data_size_qubits):
                data_size += (2**place) * sample[data_size_qubit]
            data_sizes.append(data_size)

            queue_size = 0
            for place, queue_size_qubit in enumerate(queue_size_qubits):
                queue_size += (2**place) * sample[queue_size_qubit]
            queue_sizes.append(queue_size)

            is_correct = True

            chosed_node = None
            for target_node in target_node_qubits + [empty_target_node_qubit]:
                if sample[target_node] == 1:
                    if chosed_node is not None:
                        print('Warning：不应该有两个Node都可以选')
                        is_correct = False
                    chosed_node = target_node[2]
            

            if chosed_node is None:
                print('Warning：不应该一个Node都没有选')
                is_correct = False
            
            if is_correct:
                chosed_nodes.append(chosed_node)
        
        if not is_correct:
            continue

        data_sizes = tuple(data_sizes)
        queue_sizes = tuple(queue_sizes)
        chosed_nodes = tuple(chosed_nodes)

        cases = filterCases(lambda case: case.chosed_nodes == chosed_nodes and case.data_sizes == data_sizes and case.queue_sizes == queue_sizes, all_cases)

        assert len(cases) <= 1

        if len(cases) == 0:
            case = Case()
            case.data_sizes = data_sizes
            case.queue_sizes = queue_sizes
            case.chosed_nodes = chosed_nodes
            case.energy = energy
            all_cases.append(case)
        else:
            case = cases[0]
        
        case.count += num_occurrences

        # print(data_sizes, queue_sizes, chosed_nodes, energy)
        # print(sample, energy)

    # for case in all_cases:
    #     case.ratio = case.count / read_num
    
    all_cases.sort(key=lambda elm: elm.energy)
    return all_cases


if __name__ == "__main__":
    sampler = SimulatedAnnealingSampler()
    bqm = dimod.BQM.from_qubo(Q)
    # bqm.fix_variable()
    read_num = 1000
    response = sampler.sample(bqm, num_reads=read_num)  #10000

    all_cases = analyzeCase(response)